package org.raymond.iworks.gateway.filter.support;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.symmetric.AES;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.raymond.iworks.apis.common.ErrorCode;
import org.raymond.iworks.apis.exception.BaseException;
import org.raymond.iworks.commons.utils.ErrorUtil;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.factory.rewrite.RewriteFunction;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Map;

@Slf4j
public class DecryptFunction implements RewriteFunction<String, String> {
    private final boolean encrypt;
    private final ObjectMapper objectMapper;

    public DecryptFunction(boolean encrypt, ObjectMapper objectMapper){
       this.encrypt = encrypt;
       this.objectMapper = objectMapper;
    }

    @Override
    public Publisher<String> apply(ServerWebExchange exchange, String requestString) {
        log.info("requestString:[{}]", requestString);
        GatewayContext gatewayContext = exchange.getAttribute(GatewayContext.CACHE_GATEWAY_CONTEXT);
        if (gatewayContext == null) {
            throw new BaseException(ErrorCode.ERROR, "Gateway context should not be null");
        }
        if(StringUtils.isBlank(requestString)){
            return Mono.justOrEmpty("");
        }
        if(JSONUtil.isTypeJSONArray(requestString) || JSONUtil.isTypeJSONObject(requestString)){
            throw new BaseException(ErrorCode.FORMAT_ERROR);
        }

        AES aes = gatewayContext.getAes();
        String decryptedSign = gatewayContext.getSign();
        String decrypted = null;
        if (encrypt && aes!=null) {
            decrypted = aes.decryptStr(requestString);
        } else {
            byte[] bytes = Base64.decode(requestString);
            decrypted = StrUtil.str(bytes, CharsetUtil.CHARSET_UTF_8);
        }
        if(StringUtils.isBlank(decrypted)){
            return Mono.justOrEmpty("");
        }
        log.info("Decrypted:[{}]", decrypted);
        gatewayContext.setRequestBody(decrypted);
        try {
            if(!JSONUtil.isTypeJSONObject(decrypted)){
                throw new BaseException(ErrorCode.FORMAT_ERROR);
            }
            Map requestMap = objectMapper.readValue(decrypted, Map.class);
            MultiValueMap newMap = new LinkedMultiValueMap<>();
            newMap.setAll(requestMap);
            gatewayContext.getAllRequestData().addAll(newMap);
        } catch (Exception e) {
            log.error("Read request body error: {}", ErrorUtil.getErrorLog(e));
        }
        return Mono.just(decrypted);
    }
}
